#include "WaveletFilter.h"
#include <cmath>
#include <algorithm>
#include <numeric>

// 构造函数
WaveletFilter::WaveletFilter(WaveletFunction wavelet, ThresholdRule thresholdRule)
    : m_wavelet(wavelet), m_thresholdRule(thresholdRule) {
    loadWaveletCoefficients(); // 加载小波基系数
}

// 加载小波基系数
void WaveletFilter::loadWaveletCoefficients() {
    switch (m_wavelet) {
    case DB4:
        m_lowPassDecomp = {0.4829629131445341, 0.8365163037378077, 0.2241438680420134, -0.1294095225512604};
        m_highPassDecomp = {-0.1294095225512604, -0.2241438680420134, 0.8365163037378077, -0.4829629131445341};
        m_lowPassRecon = {0.4829629131445341, 0.8365163037378077, 0.2241438680420134, -0.1294095225512604};
        m_highPassRecon = {0.1294095225512604, 0.2241438680420134, -0.8365163037378077, 0.4829629131445341};
        break;
    case SYM8:
        m_lowPassDecomp = {-0.07576571478934068, -0.02963552764595404, 0.4976186676324578, 0.8037387518052161, 0.2978577956056052, -0.09921954357693542, -0.01260396726226383, 0.03222310060407815};
        m_highPassDecomp = {-0.03222310060407815, -0.01260396726226383, 0.09921954357693542, 0.2978577956056052, -0.8037387518052161, 0.4976186676324578, 0.02963552764595404, -0.07576571478934068};
        m_lowPassRecon = {-0.07576571478934068, -0.02963552764595404, 0.4976186676324578, 0.8037387518052161, 0.2978577956056052, -0.09921954357693542, -0.01260396726226383, 0.03222310060407815};
        m_highPassRecon = {0.03222310060407815, 0.01260396726226383, -0.09921954357693542, -0.2978577956056052, 0.8037387518052161, -0.4976186676324578, -0.02963552764595404, 0.07576571478934068};
        break;
    case BIOR44:
        m_lowPassDecomp = {0.0, 0.03782845550726004, -0.023849465019556843, -0.11062440441843718, 0.37740285561283066, 0.8526986790088938, 0.37740285561283066, -0.11062440441843718, -0.023849465019556843, 0.03782845550726004};
        m_highPassDecomp = {0.0, -0.06453888262869706, 0.04068941760916406, 0.41809227322161724, -0.7884856164055829, 0.41809227322161724, 0.04068941760916406, -0.06453888262869706};
        m_lowPassRecon = {0.0, 0.06453888262869706, 0.04068941760916406, -0.41809227322161724, -0.7884856164055829, 0.41809227322161724, -0.04068941760916406, -0.06453888262869706};
        m_highPassRecon = {0.0, 0.03782845550726004, 0.023849465019556843, -0.11062440441843718, -0.37740285561283066, 0.8526986790088938, -0.37740285561283066, -0.11062440441843718, 0.023849465019556843, 0.03782845550726004};
        break;
    }
}

// 小波分解
void WaveletFilter::decompose(const QVector<double>& signal, QVector<double>& approx, QVector<double>& detail) {
    int n = signal.size();
    approx.resize(n  / 2, 0.0);
    detail.resize(n  / 2, 0.0);

    for (int i = 0; i < n / 2; ++i) {
        for (int j = 0; j < m_lowPassDecomp.size();  ++j) {
            int idx = (2 * i + j) % n; // 处理边界条件
            approx[i] += signal[idx] * m_lowPassDecomp[j];
            detail[i] += signal[idx] * m_highPassDecomp[j];
        }
    }
}

// 小波重构
QVector<double> WaveletFilter::reconstruct(const QVector<double>& approx, const QVector<double>& detail) {
    int n = approx.size()  * 2;
    QVector<double> signal(n, 0.0);

    for (int i = 0; i < n / 2; ++i) {
        for (int j = 0; j < m_lowPassRecon.size();  ++j) {
            int idx = (i - j / 2 + n / 2) % (n / 2); // 处理边界条件
            signal[2 * i] += approx[idx] * m_lowPassRecon[j];
            signal[2 * i + 1] += detail[idx] * m_highPassRecon[j];
        }
    }

    return signal;
}

// 计算阈值
double WaveletFilter::calculateThreshold(const QVector<double>& coefficients) {
    switch (m_thresholdRule) {
    case UNIVERSAL:
        return std::sqrt(2 * std::log(coefficients.size()));
    case SURE:
        return sureThreshold(coefficients);
    case HYBRID:
        return hybridThreshold(coefficients);
    }
    return 0.0;
}

// 自适应阈值
double WaveletFilter::adaptiveThreshold(const QVector<double>& coefficients) {
    double mean = std::accumulate(coefficients.begin(),  coefficients.end(),  0.0) / coefficients.size();
    double stdDev = std::sqrt(std::accumulate(coefficients.begin(),  coefficients.end(),  0.0, [mean](double sum, double val) {
                                  return sum + (val - mean) * (val - mean);
                              }) / coefficients.size());
    return stdDev * std::sqrt(2 * std::log(coefficients.size()));
}

// SURE阈值
double WaveletFilter::sureThreshold(const QVector<double>& coefficients) {
    QVector<double> sortedCoeff = coefficients;
    std::sort(sortedCoeff.begin(),  sortedCoeff.end());
    double risk = sortedCoeff.size();
    for (int i = 0; i < sortedCoeff.size();  ++i) {
        risk += sortedCoeff[i] * sortedCoeff[i] - 2 * (sortedCoeff.size()  - i) * sortedCoeff[i];
    }
    return std::sqrt(risk / sortedCoeff.size());
}

// 混合阈值
double WaveletFilter::hybridThreshold(const QVector<double>& coefficients) {
    double universal = std::sqrt(2 * std::log(coefficients.size()));
    double adaptive = adaptiveThreshold(coefficients);
    return (universal + adaptive) / 2;
}

// 滤波函数
QVector<double> WaveletFilter::filter(const QVector<double>& signal) {
    // 小波分解
    QVector<double> approx, detail;
    decompose(signal, approx, detail);

    // 阈值处理
    double threshold = calculateThreshold(detail);
    for (double& coeff : detail) {
        if (std::abs(coeff) < threshold) coeff = 0.0;
    }

    // 小波重构
    return reconstruct(approx, detail);
}
